package main

import (
    "github.com/syndtr/goleveldb/leveldb"
    "path"
    "os"
    "github.com/syndtr/goleveldb/leveldb/util"
    "github.com/boltdb/bolt"
    "encoding/json"
    "bytes"
)

type Store interface {
    Put(string, string) error
    BatchPut(map[string]string) error
    Get(string) (string, error)
    GetByPrefix(string) (map[string]string, error)
    Close() error
}

/*
for boltdb key is json ["key1","key2","key or prefix"]
*/

func NewStore(basepath, uname, tname, dbtype string) (s Store, err error) {
    switch dbtype {
    default:
        s, err = NewLevelStore(basepath, uname, tname)
    }
    return
}

type BoltStore struct {
    fn string
    db *bolt.DB
}

func parseKey(key string) (ret []string, err error) {
    err = json.Unmarshal([]byte(key), ret)
    if err != nil || len(ret) < 2 {
        return nil, InValidBoltKeyError
    }
    return
}

func (b *BoltStore) Put(k, v string) error {
    keys, err := parseKey(k)
    if err != nil {
        return err
    }
    return b.db.Update(func(tx *bolt.Tx) error {
        var b *bolt.Bucket
        var err error
        for i, key := range keys {
            switch i {
            case 0:
                if b, err = tx.CreateBucketIfNotExists([]byte(key)); err != nil {
                    return err
                }
            case len(key) - 1:
                if err = b.Put([]byte(key), []byte(v)); err != nil {
                    return err
                }
            default:
                if b, err = b.CreateBucketIfNotExists([]byte(key)); err != nil {
                    return err
                }
            }

        }
        return nil
    })
}

func (b *BoltStore) BatchPut(data map[string]string) error {
    return b.db.Batch(func(tx *bolt.Tx) error {
        for k, v := range data {
            keys, err := parseKey(k)
            if err != nil {
                return err
            }
            var b *bolt.Bucket
            for i, key := range keys {
                switch i {
                case 0:
                    if b, err = tx.CreateBucketIfNotExists([]byte(key)); err != nil {
                        return err
                    }
                case len(key) - 1:
                    if err = b.Put([]byte(key), []byte(v)); err != nil {
                        return err
                    }
                default:
                    if b, err = b.CreateBucketIfNotExists([]byte(key)); err != nil {
                        return err
                    }
                }
            }
        }
        return nil
    })
}

func (b *BoltStore) Get(k string) (string, error) {
    keys, err := parseKey(k)
    if err != nil {
        return "", err
    }
    var ret []byte
    err = b.db.View(func(tx *bolt.Tx) error {
        var b *bolt.Bucket
        var err error
        for i, key := range keys {
            switch i {
            case 0:
                if b, err = tx.CreateBucketIfNotExists([]byte(key)); err != nil {
                    return err
                }
            case len(key) - 1:
                ret = b.Get([]byte(key))
            default:
                if b, err = b.CreateBucketIfNotExists([]byte(key)); err != nil {
                    return err
                }
            }
        }
        return nil
    })
    return string(ret), err
}

func (b *BoltStore) GetByPrefix(k string) (map[string]string, error) {
    ret := make(map[string]string)
    keys, err := parseKey(k)
    if err != nil {
        return nil, err
    }
    err = b.db.View(func(tx *bolt.Tx) error {
        var b *bolt.Bucket
        var err error
        for i, key := range keys {
            switch i {
            case 0:
                if b, err = tx.CreateBucketIfNotExists([]byte(key)); err != nil {
                    return err
                }
            case len(key) - 1:
                c := b.Cursor()
                for k, v := c.Seek([]byte(key)); k != nil && bytes.HasPrefix(k, []byte(key)); k, v = c.Next() {
                    ret[string(k)] = string(v)
                }
            default:
                if b, err = b.CreateBucketIfNotExists([]byte(key)); err != nil {
                    return err
                }
            }
        }
        return nil
    })
    return ret, err
}

func (b *BoltStore) Close() error {
    return b.db.Close()
}

type LevelStore struct {
    fn string
    db *leveldb.DB
}

func NewLevelStore(basepath, uname, tname string) (Store, error) {
    dir := path.Join(basepath, uname)
    _, err := os.Stat(dir)
    if os.IsNotExist(err) {
        err = os.MkdirAll(dir, 0755)
    }
    if err != nil {
        return nil, err
    }
    fn := path.Join(dir, tname)
    db, err := leveldb.OpenFile(fn, nil)
    if err != nil {
        return nil, err
    }
    return &LevelStore{
        fn: fn, db: db,
    }, nil
}

func (l *LevelStore) Put(k, v string) error {
    return l.db.Put([]byte(k), []byte(v), nil)
}

func (l *LevelStore) GetByPrefix(prefix string) (map[string]string, error) {
    ret := make(map[string]string)
    iter := l.db.NewIterator(util.BytesPrefix([]byte(prefix)), nil)
    for iter.Next() {
        k := string(iter.Key())
        v := string(iter.Value())
        ret[k] = v
    }
    iter.Release()
    if err := iter.Error(); err != nil {
        return nil, err
    }
    return ret, nil
}

func (l *LevelStore) BatchPut(data map[string]string) error {
    batch := new(leveldb.Batch)
    for key, value := range data {
        k := []byte(key)
        v := []byte(value)
        batch.Put(k, v)
    }
    return l.db.Write(batch, nil)
}
func (l *LevelStore) Get(k string) (string, error) {
    v, err := l.db.Get([]byte(k), nil)
    return string(v), err
}
func (l *LevelStore) Close() error {
    return l.db.Close()
}
